#include "FluidSolver.h"
#include "haarwavelet.h"

#define PROJECT_ITER 40
#define DIFFUSE_ITER 10

HaarWavelet haarWvTest;

FluidSolver::FluidSolver(int w, int h, float dt) {
	mFieldA = new Field<glm::vec2>(w, h);
	mFieldB = new Field<glm::vec2>(w, h);
	mScalarFieldA = new Field<float>(w, h);
	mScalarFieldB = new Field<float>(w, h);
	SetTimeStep(dt);
	SetViscosity(1.0f);
	mHaarSelector = NULL;
}

void FluidSolver::Update( ){
	//From user forces
	CancleBoundary(*GetCurrField());
	CancleBoundary(*GetCurrScalarField());

	// XXX
	// XXX Assume same level for now.

	if (mHaarSelector!=NULL &&
		mHaarSelector->haarCellVec.size() != 0){
			int decompLevel = mHaarSelector->haarCellVec[0].level;

			haarWvTest.SetData(*(GetCurrField()->GetRawData()), GetCurrField()->GetDimX());
			haarWvTest.Decompose(decompLevel);
			HaarDecomp * decomp = haarWvTest.GetDecomp();

			for (int i = 0; i < mHaarSelector->haarCellVec.size(); ++i){
				int cellX = mHaarSelector->haarCellVec[i].cellX;
				int cellY = mHaarSelector->haarCellVec[i].cellY;
				decomp->mData[cellY*mHaarSelector->dimX + cellX] = mHaarSelector->haarCellVec[i].vel;
			}

			// Finish decomposition
			haarWvTest.Decompose(*decomp);
			haarWvTest.Reconstruct();


		int dim = GetCurrField()->GetDimX();
		Field<glm::vec2> smoothedVersion(GetCurrField()->GetDimX(), GetCurrField()->GetDimY()); //(*currField);
		for (int i = 0; i < dim; ++i){
			for (int j = 0; j < dim; ++j){
				(*GetCurrField())(i,j) = (*haarWvTest.GetData())[j*dim+i];
				//smoothedVersion(i, j) = (*haarWvTest.GetData())[j*dim+i];
			}
		}

	}


	

	//for (int i = 0; i < this->GetCurrField()->GetDimX()/2; ++i){
	//	for (int j = 0; j < this->GetCurrField()->GetDimY()/2; ++j){
	//		(*GetCurrField())(i,j) = glm::vec2(0,0);
	//	}
	//}
	//XXX

	// Velocity
	AdvectVelocity(/*&smoothedVersion*/);
	CancleBoundary(*GetCurrField());




	DiffuseVelocity();
	CancleBoundary(*GetCurrField());



	ProjectVelocity();
	CancleBoundary(*GetCurrField());

	AdvectField(/*&smoothedVersion*/);
	//CancleBoundary(*GetCurrScalarField());
}

void FluidSolver::SwapFields(){
	Field<glm::vec2> * ffp = mFieldA;
	mFieldA = mFieldB;
	mFieldB = ffp;
}

void FluidSolver::SwapScalarFields(){
	Field<float> * ffp = mScalarFieldA;
	mScalarFieldA = mScalarFieldB;
	mScalarFieldB = ffp;
}

void FluidSolver::AdvectVelocity( Field<glm::vec2> * inField ){


	// For each cell, trace a particle backwards
	glm::vec2 pos;
	glm::vec2 newPos;
	float halfCellW = 0.5f*(1.0f/(float)GetCurrField()->GetDimX());
	Field<glm::vec2> * currField = GetCurrField();
	Field<glm::vec2> * newField = GetNewField();

	Field<glm::vec2> * smoothedVersion;
	if (inField == NULL){
		smoothedVersion = (GetCurrField());
	}
	else {
		smoothedVersion = inField;
	}

	for (unsigned int i = 1; i < currField->GetDimX()-1; ++i){
		for (unsigned int j = 1; j < currField->GetDimY()-1; ++j){
			pos[0] = (float)i/currField->GetDimX() + halfCellW;
			pos[1] = (float)j/currField->GetDimX() + halfCellW;
			TraceParticle(pos, -mDt, *smoothedVersion, newPos);
			(*currField).Sample(newPos[0], newPos[1], (*newField)(i,j));
		}
	}

	SwapFields();

}

void FluidSolver::AdvectField(Field<glm::vec2> * inField){
	// For each cell, trace a particle backwards.

	glm::vec2 pos;
	glm::vec2 newPos;
	float halfCellW = 0.5f*(1.0f/(float)GetCurrScalarField()->GetDimX());
	Field<float> * currField = GetCurrScalarField();
	Field<float> * newField = GetNewScalarField();

	Field<glm::vec2> * advectField;

	if (inField == NULL){
		advectField = GetCurrField();
	}
	else {
		advectField = inField;
	}

	// Trace particle back through velocity field and 
	// use position to sample scalar field.
	for (unsigned int i = 1; i < currField->GetDimX()-1; ++i){
		for (unsigned int j = 1; j < currField->GetDimY()-1; ++j){
			pos[0] = (float)i/currField->GetDimX() + halfCellW;
			pos[1] = (float)j/currField->GetDimX() + halfCellW;
			TraceParticle(pos, -mDt, *advectField, newPos);
			currField->Sample(newPos[0], newPos[1], (*newField)(i,j));
		}
	}

	SwapScalarFields();
}

void FluidSolver::DiffuseVelocity(  ) {
	float invdim = (GetCurrField()->GetDimX() * GetCurrField()->GetDimY());
	float a = mVisc*mDt*invdim;
	float laplac[5] = { 1 + 4*a, -a, -a, -a, -a };
	PoissonSolver(laplac, DIFFUSE_ITER, *GetCurrField(), *GetNewField());

	SwapFields();

}

void FluidSolver::ProjectVelocity(  ){

	//Set up the rhs as the divergence of the current field.
	Field<glm::vec2> * currf = GetCurrField();
	Field<glm::vec2> * newf = GetNewField();
	int dimx = currf->GetDimX();
	int dimy = currf->GetDimY();
	float invdim = (dimx*dimy);
	
	Field<float>div(dimx, dimy);

	float h = dimx; //1.0f/(1.0f/dimx) == dimx
	for (unsigned int i = 1; i < dimx-1; ++i){
		for (unsigned int j = 1; j < dimy-1; ++j){
				div(i,j) = 
					((*currf)(i+1,j)[0] - (*currf)(i-1,j)[0])*0.5f
					+
					((*currf)(i,j+1)[1] - (*currf)(i,j-1)[1])*0.5f;
		}
	}

	CancleBoundary(div);

	float laplac[5] = {
		-4, 1, 1, 1, 1
	};

	// New field now contains the value of q. We subtract gradient.
	Field<float> q(div.GetDimX(), div.GetDimY());
	this->PoissonSolverScalar(laplac, dimx, dimy, PROJECT_ITER, div, q);

	// Subtract gradient of q.
	for (unsigned int i = 1; i < dimx-1; ++i){
		for (unsigned int j = 1; j < dimy-1; ++j){
			(*newf)(i, j)[0] = (*currf)(i,j)[0] - 
						(q(i+1,j) - q(i-1,j))*0.5f;
			(*newf)(i, j)[1] = (*currf)(i,j)[1] - 
						(q(i,j+1) - q(i,j-1))*0.5f;
		}
	}

	CancleBoundary(*newf);

	SwapFields();

}

void FluidSolver::SetViscosity( float visc ){
	mVisc = visc;
}

void FluidSolver::SetTimeStep(float dt ){
	mDt = dt;
}

Field<glm::vec2> * FluidSolver::GetCurrField () {
	return mFieldA;
}

Field<glm::vec2> * FluidSolver::GetNewField() {
	return mFieldB;
}

Field<float> * FluidSolver::GetCurrScalarField () {
	return mScalarFieldA;
}

Field<float> * FluidSolver::GetNewScalarField() {
	return mScalarFieldB;
}


void FluidSolver::PoissonSolver( 
	float * laplacA, 
	unsigned int iters, 
	Field<glm::vec2> & b, 
	Field<glm::vec2> & outField ){
	
	outField = b;
	Field<glm::vec2> tmpField(outField.GetDimX(), outField.GetDimY());

	Field<glm::vec2> * f1 = &outField;
	Field<glm::vec2> * f2 = &tmpField;

	for (int i = 0; i < iters; ++i){
	
		for (unsigned int u = 1; u < b.GetDimX()-1; ++u){
			for (unsigned int v = 1; v < b.GetDimY()-1; ++v){
			
				//for (unsigned int c = 0; c < 2; ++c){
				//	targetV[c] = (1.0f/laplacA[0]) * (currV[c] - 
				//		(laplacA[1]*(*f1)(u-1, v)[c] +
				//		 laplacA[2]*(*f1)(u+1, v)[c] +
				//		 laplacA[3]*(*f1)(u, v-1)[c] +
				//		 laplacA[4]*(*f1)(u, v+1)[c]));
				//}
				

				(*f2)(u, v) = (1.0f/laplacA[0]) * (b(u,v) - 
					(laplacA[1]*(*f1)(u-1, v) +
						laplacA[2]*(*f1)(u+1, v) +
						laplacA[3]*(*f1)(u, v-1) +
						laplacA[4]*(*f1)(u, v+1)));
			}
		}
	
		// Swap the fields.
		Field<glm::vec2> * tmp = f1;
		f1 = f2;
		f2 = tmp;

		CancleBoundary((*f1));
	}
	
	// f1 has the last solution.
	outField = *f1;
}

void FluidSolver::PoissonSolverScalar(float * laplacA,
	unsigned int w, 
	unsigned int h,
	unsigned int iters,
	const Field<float> & b,
	Field<float> & out) 
{
	out = b;
	Field<float> tmpField(b.GetDimX(), b.GetDimY());

	Field<float> * f1 = &out;
	Field<float> * f2 = &tmpField;

	for (int i = 0; i < iters; ++i){
	
		for (unsigned int u = 1; u < w-1; ++u){
			for (unsigned int v = 1; v < h-1; ++v){
				const float * currS = &b(u,v);
				float * targetS = &((*f2)(u,v));

					*targetS = (1.0f/laplacA[0]) * (*currS - 
						(laplacA[1]*((*f1)(u-1,v)) +
						 laplacA[2]*((*f1)(u+1,v)) +
						 laplacA[3]*((*f1)(u,v-1)) +
						 laplacA[4]*((*f1)(u,v+1))));

			}
		}
	
		// Swap the fields.
		Field<float> * tmp = f1;
		f1 = f2;
		f2 = tmp;

		CancleBoundary((*f1));
	}
	
	// f1 has the last solution.
	if (&out != f1){
		out = *f1;
	}
}
//
//void FluidSolver::PoissonSolverScalar( float * laplacA, 
//	unsigned int w, 
//	unsigned int h, 
//	unsigned int iters, 
//	std::vector<float> & b, 
//	std::vector<float> & out ){
//	
//	out = b;
//	std::vector<float> tmpField(b.size(), 0);
//
//	std::vector<float> * f1 = &out;
//	std::vector<float> * f2 = &tmpField;
//
//	for (int i = 0; i < iters; ++i){
//	
//		for (unsigned int u = 0; u < w; ++u){
//			for (unsigned int v = 0; v < h; ++v){
//				float * currS = &b[v*w + u];
//				float * targetS = &((*f2)[v*w+u]);
//
//					*targetS = (1.0f/laplacA[0]) * (*currS - 
//						(laplacA[1]*((*f1)[v*w+( (u-1) < 0 ? 0 : u-1 )]) +
//						 laplacA[2]*((*f1)[v*w+((u+1) >= w ? w-1 : u+1)]) +
//						 laplacA[3]*((*f1)[((v-1) < 0 ? 0 : v-1)*w+u]) +
//						 laplacA[4]*((*f1)[((v+1) >= h ? h-1 : v+1)*w+u]) ) );
//
//			}
//		}
//	
//		// Swap the fields.
//		std::vector<float> * tmp = f1;
//		f1 = f2;
//		f2 = tmp;
//
//	}
//	
//	// f1 has the last solution.
//	if (&out != f1){
//		out = *f1;
//	}
//}

void FluidSolver::CancleBoundary( Field<glm::vec2> & inField){
	Field<glm::vec2> * currField = &inField;
	int dimX = currField->GetDimX(),
		dimY = currField->GetDimY();
	for (int i = 1; i < dimX-1; ++i){
		(*currField)(i, 0)[1] = -(*currField)(i,1)[1]; //cancle out y
		(*currField)(i, dimY-1)[1] = -(*currField)(i,dimY-2)[1];
		(*currField)(0, i)[0] = -(*currField)(1,i)[0]; //cancle out x
		(*currField)(dimX-1, i)[0] = -(*currField)(dimX-2,i)[0];
	}
	(*currField)(0,0) = ((*currField)(1,0)+(*currField)(0,1))*0.5f;
	(*currField)(dimX-1,0) = ((*currField)(dimX-2,0)+(*currField)(dimX-1, 1))*0.5f;
	(*currField)(0,dimY-1) = ((*currField)(1,dimY-1)+(*currField)(0,dimY-2))*0.5f;
	(*currField)(dimX-1,dimY-1) = ((*currField)(dimX-2,dimY-1)+(*currField)(dimX-1,dimY-2))*0.5f;
}

void FluidSolver::CancleBoundary( Field<float> & inVec){
	int dim = inVec.GetDimX();
	for (int i = 0; i < dim-1; ++i){
		inVec(i, 0) = inVec(i,1); //cancle out y
		inVec(i, dim-1) = inVec(i,dim-2);
		inVec(0, i) = inVec(1,i); //cancle out x
		inVec(dim-1, i) = inVec(dim-2,i);
	}

	//for (int i = 0; i < dim; ++i){
	//	inVec(i, 0) = 0; //cancle out y
	//	inVec(i, dim-1) = 0;
	//	inVec(0, i) = 0; //cancle out x
	//	inVec(dim-1, i) = 0;
	//}


	inVec(0,0)=(inVec(1,0)+inVec(0,1))*0.5f;
	inVec(dim-1,0)=(inVec(dim-2,0)+inVec(dim-1,1))*0.5f;
	inVec(0,dim-1) = (inVec(1,dim-1)+inVec(0,dim-2))*0.5f;
	inVec(dim-1,dim-1) = (inVec(dim-2,dim-1)+inVec(dim-1,dim-2))*0.5f;
}


void FluidSolver::SetInterpMode( InterpMode mode ) {
	switch(mode){
	case BILIN:
        
            
            mFieldA->SetInterpMode( Field<glm::vec2>::BILIN );
		mFieldB->SetInterpMode(Field<glm::vec2>::BILIN);
		mScalarFieldA->SetInterpMode(Field<float>::BILIN);
		mScalarFieldB->SetInterpMode(Field<float>::BILIN);
		break;
	case BICUBIC:
		mFieldA->SetInterpMode(Field<glm::vec2>::BICUBIC);
		mFieldB->SetInterpMode(Field<glm::vec2>::BICUBIC);
		mScalarFieldA->SetInterpMode(Field<float>::BICUBIC);
		mScalarFieldB->SetInterpMode(Field<float>::BICUBIC);
		break;
	default:
		break;
	}
}